from mpl_toolkits import mplot3d
import numpy as np
import matplotlib.pyplot as plt

fig = plt.figure()
ax = plt.axes()


file="pydata.txt"

mnid = np.array([int(x.split(' ')[0]) for x in open(file).readlines()])
kp_x = np.array([float(x.split(' ')[1]) for x in open(file).readlines()])
kp_y = np.array([float(x.split(' ')[2]) for x in open(file).readlines()])
mp_x = np.array([float(x.split(' ')[3]) for x in open(file).readlines()])
mp_y = np.array([float(x.split(' ')[4]) for x in open(file).readlines()])
mp_z = np.array([float(x.split(' ')[5]) for x in open(file).readlines()])
ct_x = np.array([float(x.split(' ')[6]) for x in open(file).readlines()])
ct_y = np.array([float(x.split(' ')[7]) for x in open(file).readlines()])
ct_z = np.array([float(x.split(' ')[8]) for x in open(file).readlines()])

alldata = np.stack([mnid, kp_x, kp_y,mp_x,mp_y,mp_z,ct_x,ct_y,ct_z],axis=1)

frameWiseData = np.split(alldata, np.where(np.diff(alldata[:,0]))[0]+1)

for frame in frameWiseData:
	id = str(int(frame[0][0]))
	dist = np.linalg.norm(frame[:,3:6] - frame[:,6:9], axis=1) 
	dist = (dist*100) / np.linalg.norm(dist) # normalize to give uniform color shade

	ax.scatter(frame[:,1],frame[:,2],s=20,c=dist, marker = 'o', cmap = "hot" ); # https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html
	ax.axis('off')

	plt.savefig("data/frame"+id+'_.png',bbox_inches='tight') 
	ax.clear()